package top.yudoge.dogeeditor.llm.task

import dev.langchain4j.data.embedding.Embedding
import dev.langchain4j.data.segment.TextSegment
import dev.langchain4j.model.embedding.EmbeddingModel
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import lombok.extern.slf4j.Slf4j
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.scheduling.annotation.Scheduled
import org.springframework.stereotype.Component
import top.yudoge.dogeeditor.db.api.DataSourceMetaDataApi
import top.yudoge.dogeeditor.db.objects.Schema
import top.yudoge.dogeeditor.llm.embedding.EmbeddingStoreHolder

@Slf4j
@Component
class TableEmbeddingTask {

    @Autowired
    private lateinit var metaDataApi: DataSourceMetaDataApi

    @Autowired
    private lateinit var embeddingStoreHolder: EmbeddingStoreHolder

    @Autowired
    private lateinit var embeddingModel: EmbeddingModel


    @Scheduled(fixedRate = 1000 * 60 * 60) // 1 hour interval
    fun doTask() = runBlocking{
        metaDataApi.getDataSources().forEach { dataSource ->
            metaDataApi.getDatabases(dataSource.identifier).forEach { database ->
                metaDataApi.getSchemas(dataSource.identifier, database.name).forEach { schema ->
                    launch {
                        getTables(schema)
                    }
                }
            }
        }
    }

    fun getTables(schema: Schema) = runBlocking {

        val tables = metaDataApi.getTables(schema.dataSourceIdentifier, schema.databaseName, schema.name)

        val embeddingStore = embeddingStoreHolder.getStore(schema.dataSourceIdentifier, schema.databaseName, schema.name)

        tables.forEach { table ->

            val sb = StringBuilder("Create table ${table.name} (\n")
            metaDataApi.getColumns(schema.dataSourceIdentifier, schema.databaseName, schema.name, table.name).forEach { column ->
                if (column.name !in COL_BLACK_LIST) {
                    sb.append(column.name).append(" ").append(column.remark).append(",\n")
                }
            }
            sb.append(") COMMENT '${table.remark};")

            val tableName = table.name
            val tableDDL = TextSegment.from(sb.toString())

            val embeddingId = embeddingStoreHolder.genId(table.dataSourceIdentifier, schema.databaseName, schema.name, tableName)

            if (embeddingStore.hasModifications(embeddingId, tableDDL)) {
                println("Embedding for table $tableName not found, calculating...")
                var embed = embedding(tableDDL)
                embeddingStore.add(embed, tableDDL, embeddingId)
            } else {
                /* ignore. */
            }

        }

    }

    fun embedding(text: TextSegment): Embedding {
        return embeddingModel.embed(text)!!.content()
    }

    fun embedding(text: String): Embedding {
        return this.embedding(TextSegment.from(text))
    }

    companion object {
        private val COL_BLACK_LIST = setOf("id", "create_time", "update_time", "revisor", "revise_time", "is_deleted", "deleted", "delete_status", "creator", "update_user", "tenant_id")
    }


}